import triton
import itertools
import torch
import torch_npu
import triton.runtime.driver as driver


def get_autotune_config(
    multibuffer_list: list = [False],
    unit_flag_list: list = [False],
    limit_auto_multi_buffer_only_for_local_buffer_list: list = [False],
    limit_auto_multi_buffer_of_local_buffer_list: list = ["no-l0c"],
    set_workspace_multibuffer_list: list = [2, 4],
    enable_hivm_auto_cv_balance_list: list = [True],
    tile_mix_vector_loop_num_list: list = [2, 4],
    tile_mix_cube_loop_num_list: list = [2, 4],
):
    configs = []
    for (
        multibuffer,
        unit_flag,
        limit_auto_multi_buffer_only_for_local_buffer,
        limit_auto_multi_buffer_of_local_buffer,
    ) in itertools.product(
        multibuffer_list,
        unit_flag_list,
        limit_auto_multi_buffer_only_for_local_buffer_list,
        limit_auto_multi_buffer_of_local_buffer_list,
    ):

        if limit_auto_multi_buffer_only_for_local_buffer:
            configs.append(
                triton.Config(
                    {},
                    multibuffer=multibuffer,
                    unit_flag=unit_flag,
                    limit_auto_multi_buffer_only_for_local_buffer=limit_auto_multi_buffer_only_for_local_buffer,
                    limit_auto_multi_buffer_of_local_buffer=limit_auto_multi_buffer_of_local_buffer,
                )
            )
        else:
            for (
                set_workspace_multibuffer,
                enable_hivm_auto_cv_balance,
                tile_mix_vector_loop,
                tile_mix_cube_loop,
            ) in itertools.product(
                set_workspace_multibuffer_list,
                enable_hivm_auto_cv_balance_list,
                tile_mix_vector_loop_num_list,
                tile_mix_cube_loop_num_list,
            ):
                configs.append(
                    triton.Config(
                        {},
                        multibuffer=multibuffer,
                        unit_flag=unit_flag,
                        limit_auto_multi_buffer_only_for_local_buffer=limit_auto_multi_buffer_only_for_local_buffer,
                        limit_auto_multi_buffer_of_local_buffer=limit_auto_multi_buffer_of_local_buffer,
                        set_workspace_multibuffer=set_workspace_multibuffer,
                        enable_hivm_auto_cv_balance=enable_hivm_auto_cv_balance,
                        tile_mix_vector_loop=tile_mix_vector_loop,
                        tile_mix_cube_loop=tile_mix_cube_loop,
                    )
                )
    return configs


def get_npu_properties():
    device = torch.npu.current_device()
    return driver.active.utils.get_device_properties(device)